import numpy as np
from astropy.io import fits
import argparse
import os

def read_binary_file(filename, dtype=np.uint16):
    """从二进制文件中读取数据并转换为 NumPy 数组"""
    with open(filename, 'rb') as file:
        data = np.fromfile(file, dtype=dtype)
    return data

def split_data_into_columns(data, num_columns=16):
    """将数据按列分割成多个子数组"""
    if len(data) % num_columns != 0:
        raise ValueError("数据长度必须是列数的整数倍")
    columns = []
    for i in range(num_columns):
        column = data[i::num_columns]
        columns.append(column.reshape(1050, 1100))  # 重塑为1050×1100图像
    return columns

# def average_adjacent_columns(columns):
#     """将相邻的两列图像进行平均处理（使用完整的16列）"""
#     if len(columns) != 16:
#         raise ValueError("必须提供16列图像数据进行处理")
    
#     averaged_columns = []
    
#     # 处理第一行的8列(索引0-7)
#     for i in range(0, 8, 2):
#         avg = (columns[i] + columns[i+1]) / 2
#         averaged_columns.append(avg)
    
#     # 处理第二行的8列(索引8-15)
#     for i in range(8, 16, 2):
#         avg = (columns[i] + columns[i+1]) / 2
#         averaged_columns.append(avg)
    
#     return averaged_columns

def concatenate_columns(columns, num_rows, num_columns_per_row):
    """将列数据拼接为二维数组"""
    if len(columns) != num_rows * num_columns_per_row:
        raise ValueError("列的数量必须是每行列数的整数倍")
    rows = []
    for i in range(num_rows):
        start = i * num_columns_per_row
        end = start + num_columns_per_row
        row = np.hstack(columns[start:end])  # 水平拼接每行的列
        rows.append(row)
    combined_image = np.vstack(rows)  # 垂直拼接各行
    return combined_image

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="处理二进制文件并生成筛选后的原始图像和平均后的FITS图像")
    parser.add_argument("bin_file", help="输入的二进制文件路径")
    parser.add_argument("output_dir", help="输出目录")
    parser.add_argument("output_name", nargs="?", default="mosaic_2x8.fits", help="输出FITS文件名（可不带扩展名，默认：mosaic_2x8.fits）")
    args = parser.parse_args()
    
    # 确保输出目录存在
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)
    
    # 1. 读取数据
    data = read_binary_file(args.bin_file)
    
    # 2. 分割数据为16列
    all_columns = split_data_into_columns(data)
    
    # 3. 生成筛选后的原始图像（只保留第1、3、5、7 9 11 13 15列）
    # selected_indices = [2, 3, 6, 7, 8, 9, 12, 13]  # 对应第1、3、5、7,9 11 13 15列（Python从0开始计数）
    selected_indices = [0, 4, 8, 12] 
    selected_columns = [all_columns[i] for i in selected_indices]
    selected_image = concatenate_columns(selected_columns, num_rows=2, num_columns_per_row=2)

     # 4. 写出FITS（允许第三个位置参数省略 .fits 后缀）
    output_name = args.output_name
    if not output_name.lower().endswith('.fits'):
        output_name = output_name + '.fits'
    output_path = os.path.join(args.output_dir, output_name)
    fits.PrimaryHDU(data=selected_image).writeto(output_path, overwrite=True)
    print(f"2×2拼接的FITS文件已生成：{output_path}")